use std::fs::File;
use std::io::BufReader;
use std::f32::consts::PI;

use crate::jpeg_meta::*;
use crate::jpeg_reader::*;
use crate::jpeg_utils::*;
use crate::jpeg_image::*;

const ZigZag: [[usize; 8]; 8] = [
    [ 0,  1,  5,  6, 14, 15, 27, 28 ],
    [ 2,  4,  7, 13, 16, 26, 29, 42 ],
    [ 3,  8, 12, 17, 25, 30, 41, 43 ],
    [ 9, 11, 18, 24, 31, 40, 44, 53 ],
    [ 10, 19, 23, 32, 39, 45, 52, 54 ],
    [ 20, 22, 33, 38, 46, 51, 55, 60 ],
    [ 21, 34, 37, 47, 50, 56, 59, 61 ],
    [ 35, 36, 48, 49, 57, 58, 62, 63 ]
];

struct JpegMCUData<'a> {
    mcu: MCUInfo,
    jpeg_meta: &'a JPEGMetaData
}

impl<'a> JpegMCUData<'a> {

    pub fn new(mcu: MCUInfo, jpeg_meta: &'a JPEGMetaData) -> Self {
        JpegMCUData {
            mcu,
            jpeg_meta,
        }
    }

    pub fn decode(&mut self) {
        self.dequantize();
        self.zigzag();
        self.idct();
    }

    pub fn dequantize(&mut self) {
        let quantable = self.jpeg_meta.quantable;
        let sof = &self.jpeg_meta.sof;
        for id in 0..sof.num_com as usize{
            let com = &sof.coms[id];
            for r in 0..com.v_sample_factor as usize {
                for c in 0..com.h_sample_factor as usize {
                    self.mcu[id][r][c] = jpeg_utils::multi(self.mcu[id][r][c], quantable[com.quantable_id as usize]);
                }
            }
        }
    }

    pub fn zigzag(&mut self) {
        let sof = &self.jpeg_meta.sof;
        for id in 0..sof.num_com as usize {
            let com = &sof.coms[id];
            for r in 0..com.v_sample_factor as usize {
                for c in 0..com.h_sample_factor as usize {

                    let mut tmp: [[f32;8];8] = Default::default();
                    for i in 0..8 {
                        for j in 0..8 {
                            tmp[i][j] = self.mcu[id][r][c][ZigZag[i][j] / 8][ZigZag[i][j] % 8];
                        }
                    }
                    self.mcu[id][r][c] = tmp;

                }
            }
        }
    }

    pub fn idct(&mut self) {

        fn cc(i: usize, j: usize) -> f32 {
            if i == 0 && j == 0 {
                return 1.0 / 2.0;
            } else if i == 0 || j == 0 {
                return 1.0 / (2.0 as f32).sqrt();
            } else {
                return 1.0;
            }
        }

        let sof = &self.jpeg_meta.sof;
        for id in 0..sof.num_com as usize {
            let com = &sof.coms[id];
            for r in 0..com.v_sample_factor as usize {
                for c in 0..com.h_sample_factor as usize {

                    let mut f = [[0f32;8];8];
                    for i in 0..8 {
                        for j in 0..8 {
                            for u in 0..8 {
                                for v in 0..8 {
                                    let cosu = ((i as f32 + 0.5) * PI * u as f32 / 8.0).cos();
                                    let cosv = ((j as f32 + 0.5) * PI * v as f32 / 8.0).cos();
                                    f[i][j] += cc(u,v)*self.mcu[id][r][c][u][v] * cosu * cosv;
                                }
                            }
                            f[i][j] /= 4.0;
                        }
                    }

                    self.mcu[id][r][c] = f;
                }
            }
        }
    }

    pub fn to_rgb(&mut self) -> Vec<Vec<Color>> {
        self.decode();

        let sof = &self.jpeg_meta.sof;
        let coms = &sof.coms;
        let max_v_sample = sof.v_max_samples;
        let max_h_smaple = sof.h_max_samples;
        let mcu_width = (max_h_smaple * 8) as usize;
        let mcu_height = (max_v_sample * 8) as usize;

        fn to_u8(x:f32) -> u8 {
            if x >= 255.0 {
                return 255;
            } else if x <= 0.0 {
                return 0;
            } else {
                return x.round() as u8;
            }
        }

        let mut ret = vec![vec![Color::new(); mcu_width]; mcu_height];
        for h in 0..mcu_height {
            for w in 0..mcu_width {
                // 获取 Y Cb Cr 三个颜色分量对应位置的采样值
                let mut YCbCr = [0.0; 3];
                for id in 0..3 {
                    let vh = (h as u8 * coms[id].v_sample_factor / max_v_sample) as usize;
                    let vw = (w as u8 * coms[id].h_sample_factor / max_h_smaple) as usize;
                    YCbCr[id] = self.mcu[id][vh/8][vw/8][vh%8][vw%8];
                }
                let (Y, Cb, Cr) = (YCbCr[0], YCbCr[1], YCbCr[2]);
                let R = to_u8(Y + 1.402 * Cr + 128.0);
                let G = to_u8(Y - 0.34414 * Cb - 0.71414 * Cr + 128.0);
                let B = to_u8(Y + 1.772 * Cb + 128.0);

                ret[h][w] = Color::RGB(R, G, B);
            }
        }

        return ret;
    }
}

pub mod jpeg_decode {
    use super::*;

    pub fn data_decode(reader: &mut BufReader<File>) -> Image {
        let (jpeg_meta, mcus, ) = jpeg_read::data_read(reader);

        let sof = &jpeg_meta.sof;

        let mcu_width = 8 * sof.h_max_samples as usize;
        let mcu_height = 8 * sof.v_max_samples as usize;

        let h_mcu_num = (sof.width as usize - 1 ) / mcu_width + 1;
        let v_mcu_num = (sof.height as usize - 1 ) / mcu_height + 1;

        let mut image = Image::new( (mcu_width * h_mcu_num) as u32, (mcu_height * v_mcu_num) as u32);

        for r in 0..v_mcu_num as usize {
            for c in 0..h_mcu_num as usize {
                let mcu = mcus[r][c].clone();
                let mcu_rgb = JpegMCUData::new(mcu, &jpeg_meta).to_rgb();
                for h in 0..mcu_height {
                    for w in 0..mcu_width {
                        image.pixels[ r * mcu_height + h ][ c * mcu_width + w] = mcu_rgb[h][w];
                    }
                }
            }
        }

        return image;
    }
}
